#!/usr/bin/env python3
#-*- coding:utf-8 -*-

import os
import sys
import re
import argparse
import subprocess
import threading
import time

from transitions import Machine, State

_fpath = os.path.dirname(os.path.realpath(__file__))
_eeldir = _fpath

import logging
_log = logging

def_fmt = '[%(asctime)s][%(levelname)-8s] %(message)s'
def_datefmt = '%Y%m%d %H:%M:%S'

def setup_logging(args):
    # logfile, log_level=logging.DEBUG, log_format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'):
    # 'global _log' must in function context
    global _log
    _log = logging.getLogger("autotest")
    _log.propagate = False
    _log.setLevel(logging.DEBUG)

    loglvl_sh = eval(f'logging.{args.loglevel}')

    logfmt_sh = logging.Formatter(def_fmt, datefmt=None, style='%')
    # take effective when datefmt is None
    logfmt_sh.default_time_format = def_datefmt
    logfmt_sh.default_msec_format = "%s.%03ld"

    # StreamHandler for terminal output
    sh = logging.StreamHandler()
    sh.setLevel(loglvl_sh)
    sh.setFormatter(logfmt_sh)
    _log.addHandler(sh)

    datatime = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    logfile = os.path.join(_eeldir, "testlog", f"test_{datatime}.log")
    if not os.path.exists(os.path.dirname(logfile)):
        os.makedirs(os.path.dirname(logfile))
    loglvl_fh = logging.DEBUG
    logfmt_fh = logging.Formatter(def_fmt, datefmt=None, style='%')
    # take effective when datefmt is None
    logfmt_fh.default_time_format = def_datefmt
    logfmt_fh.default_msec_format = "%s.%03ld"
    # FileHandler for file output
    fh = logging.FileHandler(logfile)
    fh.setLevel(loglvl_fh)
    fh.setFormatter(logfmt_fh)
    _log.addHandler(fh)

    setattr(_log, 'stdoutlog', eval(f'_log.{args.stdout_level.lower()}'))

fail_patterns = (
    re.compile('^make: .* Error'),
    re.compile('^make: .* Stop'),
)

success_patterns_ub = (re.compile('Hit any key to stop'), )

success_patterns_kl = (re.compile('^.*end Kernel panic.*VFS: Unable to mount root fs', re.IGNORECASE),)   # arm64_tiny
success_patterns_br = (
    # re.compile('^buildroot login:'),      # can not use this
    re.compile('^Welcome to Buildroot'),
)

class Config:
    """Configuration management class"""
    def __init__(self):
        # some config can not run by qemu
        self.presets_targets = {
            "./presets/arm64_tiny.config.mk"    : ('qemu-uboot', 'qemu-kernel', 'qemu'),
            "./presets/arm64.config.mk"         : ('qemu-uboot', 'qemu-kernel', 'qemu'),
            "./presets/arm64.tee.config.mk"     : ('qemu-uboot', 'qemu'),
            "./presets/arm64.perf.config.mk"    : ('qemu-uboot', 'qemu-kernel', 'qemu'),
            "./presets/arm.config.mk"           : ('qemu-uboot', 'qemu-kernel', 'qemu'),
            "./presets/arm.tee.config.mk"       : ('qemu-uboot', 'qemu'),
            "./presets/riscv64.config.mk"       : ('qemu-uboot', 'qemu-kernel', 'qemu'),
            "./presets/riscv32.config.mk"       : ('qemu-uboot', 'qemu-kernel', 'qemu'),
        }

        self.targets_presets = {
            'qemu-uboot': (
                "./presets/arm64_tiny.config.mk",
                "./presets/arm64.config.mk",
                "./presets/arm64.tee.config.mk",
                "./presets/arm64.perf.config.mk",
                "./presets/arm.config.mk",
                "./presets/arm.tee.config.mk",
                "./presets/riscv64.config.mk",
                "./presets/riscv32.config.mk",
            ),
            'qemu-kernel' : (
                "./presets/arm64_tiny.config.mk",
                "./presets/arm64.config.mk",
                "./presets/arm64.perf.config.mk",
                "./presets/arm.config.mk",
                "./presets/riscv64.config.mk",
                "./presets/riscv32.config.mk",
            ),
            'qemu': (
                "./presets/arm64_tiny.config.mk",
                "./presets/arm64.config.mk",
                "./presets/arm64.tee.config.mk",
                "./presets/arm64.perf.config.mk",
                "./presets/arm.config.mk",
                "./presets/arm.tee.config.mk",
                "./presets/riscv64.config.mk",
                "./presets/riscv32.config.mk",
            ),
        }

        self.common_configs = {
            'qemu-uboot': {
                'triggers': {
                    'fail': fail_patterns,
                    'success': success_patterns_ub,
                },
            },
            'qemu-kernel': {
                'triggers': {
                    'fail': fail_patterns,
                    'success': success_patterns_kl + success_patterns_br,
                }
            },
            'qemu': {       # for tee
                'triggers': {
                    'fail': fail_patterns,
                    'success': success_patterns_kl + success_patterns_br,
                }
            }
        }

    def get_timeout(self, target: str) -> float:
        timeout_map = {
            'qemu-uboot': 120,
            'qemu-kernel': 2 * 60 * 60,
            'qemu': 2 * 60 * 60
        }
        return timeout_map.get(target, 60)

    def get_triggers(self, target:str, cfg:str) -> dict:
        if 'triggers' in cfg:
            return cfg['triggers']

        return self.common_configs[target]['triggers']


class EMatter(object):
    def __init__(self):
        self._process = None
        states = (
            State('start', on_exit=['run_start']),
            'working',
            State('successed', on_enter=['do_successed']),
            State('failed', on_enter=['do_failed']),
        )
        transitions = (
            {'trigger': 'run',      'source': 'start',   'dest': 'working', 'before': ['time_reset'],},
            {'trigger': 'success',  'source': 'working', 'dest': 'successed',},
            {'trigger': 'fail',     'source': 'working', 'dest': 'failed',},
            {'trigger': 'reset',    'source': '*',       'dest': 'start',}
        )
        self._machine = Machine(model=self,
                states=states, transitions=transitions,
                initial='start')
        self._lock = threading.Lock()
        self._run_tstart = 0


    def get_state(self):
        return self._machine.get_model_state(self)

    def run_start(self, *args):
        self._run_tstart = time.time()

    def time_reset(self, *args):
        self._tstart = time.time()

    def timeout(self, seconds:float):
        if seconds < 0:
            return False
        return time.time() - self._tstart > seconds

    def _qemu_terminate(self):
        """Safely terminate QEMU process with error handling"""
        try:
            if self._process and self._process.stdin:
                for ch in (chr(1), 'x'):
                    self._process.stdin.write(ch)
                    self._process.stdin.flush()
        except Exception as e:
            _log.error(f"Failed to terminate QEMU: {e}")
        # finally:
        #     if self._process:
        #         self._process.terminate()

    def do_successed(self, line):
        self.time_reset()
        self._qemu_terminate()
        mesg = f"Success: {line}"
        _log.info(mesg)

    def do_failed(self, line):
        self._qemu_terminate()
        _log.error(f"Failed: {line}")

    def _process_line(self, triggers:dict, line:str):
        self._lock.acquire()
        for triggername, patterns in triggers.items():
            for pattern in patterns:
                if pattern.match(line):
                    self.trigger(triggername, line)
                    break
        self._lock.release()

    def _process_stdout(self, triggers:dict):
        for line in self._process.stdout:
            _log.stdoutlog(line.strip('\n\r'))
            self._process_line(triggers, line)

    def _process_stderr(self, triggers:dict):
        for line in self._process.stderr:
            # if not line:
            #     print("-------")  # would not print
            #     continue
            line = line.strip('\n\r')
            _log.error(f"stderr: {line}")
            self._process_line(triggers, line)

    def _check_timetout(self, seconds:float, interval:float=0.1):
        while self._process.poll() is None:
            if self.timeout(seconds):
                _log.error(f"Timeout: {self._tstart}")
                self._lock.acquire()
                self.trigger('fail', 'timeout')
                self._lock.release()
                break
            time.sleep(interval)

    def process(self, config:dict):
        cmd = config.get('command')
        self._process = subprocess.Popen(cmd, cwd=config.get("workdir"),
                    shell=True, text=True,
                    stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        self.trigger('run')

        threads = []
        th = threading.Thread(target=self._process_stderr, args=(config.get("triggers"),))
        threads.append(th)
        th = threading.Thread(target=self._process_stdout, args=(config.get("triggers"),))
        threads.append(th)
        th = threading.Thread(target=self._check_timetout, args=(config.get("timeout"), ))
        threads.append(th)

        for thread in threads:
            thread.daemon = True
            thread.start()

        try:
            _log.info(f"Processing: {cmd}")
            self._process.wait()
        except KeyboardInterrupt:
            _log.info(f"\nKeyboard Interrupt")
            self._process.terminate()
        finally:
            for thread in threads:
                thread.join()

class Tester(object):

    def __init__(self, config=Config()):
        self._config = config
        self._ematter = EMatter()
        self._workdir = _eeldir

    def _get_command(self, args, target:str, configfile:str=''):
        command = f'make {target}'
        if args.output:
            command += f" O={args.output}"
        if configfile:
            command += f" CONF={configfile}"
        return command

    def _get_timeout(self, args, target:str, cfg:str) -> float:
        if 'timeout' in cfg:
            return cfg["timeout"]

        # if 'submodule' in args.cleans:
        #     return 60
        if target == 'qemu-uboot':
            return 120

        if target in ('qemu-kernel', 'qemu'):
            return 2 * 60 * 60

        raise ValueError(f"Unknown target: {target}")

    def do_clean(self, args):
        targets = args.cleans

    def run_test(self, args, target: str, cfg: str) -> tuple:
        """Execute a single test with better error handling"""
        try:
            test_config = {
                'workdir': self._workdir,
                'command': self._get_command(args, target, cfg),
                'timeout': self._config.get_timeout(target),
                'triggers': self._config.get_triggers(target, cfg),
            }

            if args.cleans:
                self.do_clean(args)

            self._ematter.process(test_config)
            result = self._ematter.get_state()
            self._ematter.trigger('reset')
            return (cfg, target, result.value)

        except Exception as e:
            _log.error(f"Test failed: {e}")
            return (cfg, target, 'error')

    def iter_presets_targets(self, args, presets:list, targets:list):
        results = []
        for preset in presets:
            for target in targets:
                ret = self.run_test(args, target, preset)
                results.append(ret)
        for ret in results:
            print(ret)

    def iter_targets_presets(self, args, presets:list, targets:list):
        results = []
        for target in targets:
            for preset in presets:
                ret = self.run_test(args, target, preset)
                results.append(ret)
        for ret in results:
            print(ret)

    def do_test(self, args, presets:list, targets:list):
        if args.first == 'preset':
            return self.iter_presets_targets(args, presets, targets)
        if args.first == 'target':
            return self.iter_targets_presets(args, presets, targets)
        raise ValueError(f"Unknown first: {args.first}")


def autotest(args, config):
    tester = Tester(config=config)
    results = []
    if args.first == 'preset':
        for preset, targets in config.presets_targets.items():
            for target in targets:
                ret = tester.run_test(args, target, cfg=preset)
                results.append(ret)
    elif args.first == 'target':
        for target, presets in config.targets_presets.items():
            for preset in presets:
                ret = tester.run_test(args, target, cfg=preset)
                results.append(ret)

    for ret in results:
        print(ret)

def main(args):
    try:
        config = Config()

        if args.autotest:
            return autotest(args, config)

        if not args.presets or not args.targets:
            raise ValueError("Please specify the preset and target.")

        tester = Tester(config=config)
        tester.do_test(args, args.presets, args.targets)

    except Exception as e:
        _log.error(f"Fatal error: {e}")
        sys.exit(1)

def parseargs():
    argp = argparse.ArgumentParser(description='auto make test')

    argp.add_argument("-l", "--loglevel", type=str,
                      choices=('ERROR', 'WARNING', 'INFO', 'DEBUG'), default='INFO',
                      help="loglevel")
    argp.add_argument("--stdout-level", type=str,
                      choices=('INFO', 'DEBUG'), default='DEBUG',
                      help="loglevel of subprocess stdout")

    argp.add_argument("--autotest", action="store_true", help="auto test")

    argp.add_argument("-c", "--cleans",  type=str, nargs="+", default=[], help="clean target(s) before test, submodule/source/output/none")
    argp.add_argument("-p", "--presets", type=str, nargs='*', default=[], help="preset(s) file path")
    argp.add_argument("-t", "--targets", type=str, nargs="+", default=['none'], help="target(s) to test")
    argp.add_argument("-F", "--first", type=str, choices=('preset', 'target'), default='preset', help='first iter preset or target')
    argp.add_argument("-o", "--output", type=str, default='', help="output dir")

    args = argp.parse_args()
    return args

if __name__ == "__main__":
    args = parseargs()
    setup_logging(args)
    _log.info(f"{args}")
    main(args)
